/*
 * @Author: Zhou Zijian
 * @Date: 2023-03-01 01:46:29
 * @Last Modified by: Zhou Zijian
 * @Last Modified time: 2023-05-07 17:37:00
 */

#include "GpuX.h"
#include <fstream>
#include <memory>
#include <regex>
#include "ConfigParser.h"
#include "lib-loader/ILibLoader.h"

int GpuX::OpenclLibLoader::Open(const char *libName)
{
    int ret = OpenLib(libName, &m_handle);
    if (ret != 0) {
        return GpuX::ErrorCode::OPEN_LIB_FAILED;
    }
    GetPlatformIDs = static_cast<clGetPlatformIDsPtrFunc>(LoadFunc(m_handle, "clGetPlatformIDs"));
    GetDeviceIDs = static_cast<clGetDeviceIDsPtrFunc>(LoadFunc(m_handle, "clGetDeviceIDs"));
    CreateContext = static_cast<clCreateContextPtrFunc>(LoadFunc(m_handle, "clCreateContext"));
    CreateCommandQueue = static_cast<clCreateCommandQueuePtrFunc>(LoadFunc(m_handle, "clCreateCommandQueue"));
    CreateProgramWithSource = static_cast<clCreateProgramWithSourcePtrFunc>(LoadFunc(m_handle,
        "clCreateProgramWithSource"));
    BuildProgram = static_cast<clBuildProgramPtrFunc>(LoadFunc(m_handle, "clBuildProgram"));
    CreateKernel = static_cast<clCreateKernelPtrFunc>(LoadFunc(m_handle, "clCreateKernel"));
    ReleaseKernel = static_cast<clReleaseKernelPtrFunc>(LoadFunc(m_handle, "clReleaseKernel"));
    ReleaseProgram = static_cast<clReleaseProgramPtrFunc>(LoadFunc(m_handle, "clReleaseProgram"));
    ReleaseCommandQueue = static_cast<clReleaseCommandQueuePtrFunc>(LoadFunc(m_handle, "clReleaseCommandQueue"));
    ReleaseContext = static_cast<clReleaseContextPtrFunc>(LoadFunc(m_handle, "clReleaseContext"));
    SetKernelArg = static_cast<clSetKernelArgPtrFunc>(LoadFunc(m_handle, "clSetKernelArg"));
    EnqueueNDRangeKernel = static_cast<clEnqueueNDRangeKernelPtrFunc>(LoadFunc(m_handle, "clEnqueueNDRangeKernel"));
    Finish = static_cast<clFinishPtrFunc>(LoadFunc(m_handle, "clFinish"));
    CreateBuffer = static_cast<clCreateBufferPtrFunc>(LoadFunc(m_handle, "clCreateBuffer"));
    EnqueueReadBuffer = static_cast<clEnqueueReadBufferPtrFunc>(LoadFunc(m_handle, "clEnqueueReadBuffer"));
    ReleaseMemObject = static_cast<clReleaseMemObjectPtrFunc>(LoadFunc(m_handle, "clReleaseMemObject"));
    if (GetPlatformIDs == nullptr || GetDeviceIDs == nullptr || CreateContext == nullptr || 
        CreateCommandQueue == nullptr || CreateProgramWithSource == nullptr || BuildProgram == nullptr || 
        CreateKernel == nullptr || ReleaseKernel == nullptr || ReleaseProgram == nullptr || 
        ReleaseCommandQueue == nullptr || ReleaseContext == nullptr || SetKernelArg == nullptr || 
        EnqueueNDRangeKernel == nullptr || Finish == nullptr || CreateBuffer == nullptr ||
        EnqueueReadBuffer == nullptr || ReleaseMemObject == nullptr) {
        return GpuX::ErrorCode::LOAD_FUNC_FAILED;
    }
    return GpuX::ErrorCode::SUCCESS;
}

int GpuX::OpenclLibLoader::Close()
{
    return CloseLib(m_handle) == 0 ? GpuX::ErrorCode::SUCCESS : GpuX::ErrorCode::CLOSE_LIB_FAILED;
}

GpuX::Buffer::~Buffer()
{
    if (m_mem != 0) {
        m_cl->ReleaseMemObject(m_mem);
        m_mem = 0;
    }
}

int GpuX::Buffer::Init(void *ptr, size_t size)
{
    if (ptr == nullptr) {
        return GpuX::ErrorCode::CREATE_BUFFER_FAILED_NULLPTR;
    }
    if (size <= 0) {
        return GpuX::ErrorCode::CREATE_BUFFER_FAILED_SIZE_ZERO;
    }
    cl_int err;
    m_mem = m_cl->CreateBuffer(m_context, CL_MEM_READ_WRITE | CL_MEM_USE_HOST_PTR, size, ptr, &err);
    if (err != CL_SUCCESS) {
        return GpuX::ErrorCode::CREATE_BUFFER_FAILED;
    }
    m_size = size;
    m_ptr = ptr;
    return GpuX::ErrorCode::SUCCESS;
}

cl_mem GpuX::Buffer::GetClMem()
{
    return m_mem;
}

void *GpuX::Buffer::GetPointer()
{
    return m_ptr;
}

size_t GpuX::Buffer::GetSize()
{
    return m_size;
}

int GpuX::Buffer::Flush()
{
    cl_int err = m_cl->EnqueueReadBuffer(m_commandQueue, m_mem, CL_TRUE, 0, m_size, m_ptr, 0, nullptr, nullptr);
    if (err != CL_SUCCESS) {
        return GpuX::ErrorCode::FLUSH_BUFFER_FAILED;
    }
    return GpuX::ErrorCode::SUCCESS;
}

int GpuX::GetDeviceId()
{
    // get platform id
    cl_uint num_platform;
    m_err = m_cl.GetPlatformIDs(0, nullptr, &num_platform);
    if (m_err != CL_SUCCESS) {
        return GpuX::ErrorCode::GET_PLATFORM_ID_FAILED;
    }
    if (num_platform == 0) {
        return GpuX::ErrorCode::NO_PLATFORM;
    }
    std::unique_ptr<cl_platform_id[]> platformId(new cl_platform_id[num_platform]);
    m_err = m_cl.GetPlatformIDs(num_platform, platformId.get(), nullptr);
    if (m_err != CL_SUCCESS) {
        return GpuX::ErrorCode::GET_PLATFORM_ID_FAILED;
    }

    // get device id
    cl_uint num_device;
    for (int i = 0; i < num_platform; i++) {
        m_err = m_cl.GetDeviceIDs(platformId[i], CL_DEVICE_TYPE_GPU, 0, nullptr, &num_device);
        if (m_err != CL_SUCCESS) {
            return GpuX::ErrorCode::GET_DEVICE_ID_FAILED;
        }
        if (num_device != 0) {
            m_err = m_cl.GetDeviceIDs(platformId[i], CL_DEVICE_TYPE_GPU, 1, &m_deviceId, nullptr);
            if (m_err != CL_SUCCESS) {
                return GpuX::ErrorCode::GET_DEVICE_ID_FAILED;
            }
            break;
        }
    }
    if (num_device == 0) {
        return GpuX::ErrorCode::NO_GPU;
    }
    return GpuX::ErrorCode::SUCCESS;
}

int GpuX::GetContext()
{
    // create context
    m_context = m_cl.CreateContext(nullptr, 1, &m_deviceId, nullptr, nullptr, &m_err);
    if (m_err != CL_SUCCESS) {
        return GpuX::ErrorCode::CREATE_CONTEXT_FAILED;
    }
    return GpuX::ErrorCode::SUCCESS;
}

int GpuX::GetCommandQueue()
{
    // create command queue
    m_commandQueue = m_cl.CreateCommandQueue(m_context, m_deviceId, 0, &m_err);
    if (m_err != CL_SUCCESS) {
        return GpuX::ErrorCode::CREATE_COMMAND_QUEUE_FAILED;
    }
    return GpuX::ErrorCode::SUCCESS;
}

int GpuX::BuildProgram(const std::string &program)
{
    std::ifstream f(program, std::ifstream::binary);
    if (!f.is_open()) {
        return ErrorCode::OPEN_FILE_FAILED;
    }

    // check program
    std::regex regCl(".+.cl");
    bool isCl = std::regex_match(program, regCl);
    if (!isCl) {
        return NO_BIN_OR_CL;
    }

    f.seekg(0, f.end);
    const size_t programSize = f.tellg();
    f.seekg(0, f.beg);

    std::unique_ptr<char[]> programBuffer(new char[programSize + 1]);
    programBuffer[programSize] = '\0';
    f.read(programBuffer.get(), programSize);
    constexpr int programNum = 1;
    const char *programBuffers[programNum] = {programBuffer.get()};
    const size_t programSizes[programNum] = {programSize};

    m_program = m_cl.CreateProgramWithSource(m_context, programNum, programBuffers, programSizes, &m_err);
    if (m_err != CL_SUCCESS) {
        return GpuX::ErrorCode::CREATE_PROGRAM_FAILED;
    }

    const char options[] = "-cl-mad-enable -Werror";
    m_err = m_cl.BuildProgram(m_program, 1, &m_deviceId, options, nullptr, nullptr);
    if (m_err != CL_SUCCESS) {
        return GpuX::ErrorCode::BUILD_PROGRAM_FAILED;
    }
    return GpuX::ErrorCode::SUCCESS;
}

int GpuX::SetFunctions(const std::string &name,
    const std::string &kernelName,
    const std::vector<size_t> &globalWorkSize,
    const std::vector<size_t> &localWorkSize)
{
    cl_kernel kernel = m_cl.CreateKernel(m_program, kernelName.c_str(), &m_err);
    if (m_err != CL_SUCCESS) {
        return GpuX::ErrorCode::CREATE_KERNEL_FAILED;
    }
    m_kernels.insert(kernel);
    m_functions.insert(std::pair<std::string, cl_kernel>(name, kernel));
    m_globalWorkSize.insert(std::pair<std::string, std::vector<size_t>>(name, globalWorkSize));
    m_localWorkSize.insert(std::pair<std::string, std::vector<size_t>>(name, localWorkSize));
    return GpuX::ErrorCode::SUCCESS;
}

int GpuX::Create(const std::string &config)
{
    ConfigParser configParser;
    int ret = configParser.LoadJson(config);
    if (ret != ConfigParser::ErrorCode::SUCCESS) {
        return GpuX::ErrorCode::CONFIG_FAILED;
    }
    auto program = configParser.GetProgram();
    auto functions = configParser.GetFunctions();
    auto opencl = configParser.GetOpencl();

    ret = m_cl.Open(opencl.c_str());
    if (ret != GpuX::ErrorCode::SUCCESS) {
        return ret;
    }

    ret = GetDeviceId();
    if (ret != GpuX::ErrorCode::SUCCESS) {
        return ret;
    }

    ret = GetContext();
    if (ret != GpuX::ErrorCode::SUCCESS) {
        Destroy();
        return ret;
    }

    ret = GetCommandQueue();
    if (ret != GpuX::ErrorCode::SUCCESS) {
        Destroy();
        return ret;
    }

    ret = BuildProgram(program);
    if (ret != GpuX::ErrorCode::SUCCESS) {
        Destroy();
        return ret;
    }

    for (auto &function : functions) {
        ret = SetFunctions(function.name, function.kernel, function.globalWorkSize, function.localWorkSize);
        if (ret != GpuX::ErrorCode::SUCCESS) {
            Destroy();
            return ret;
        }
    }

    return GpuX::ErrorCode::SUCCESS;
}

void GpuX::Destroy()
{
    for (auto &kernel : m_kernels) {
        if (kernel != 0) {
            m_cl.ReleaseKernel(kernel);
        }
    }

    if (m_program != 0) {
        m_cl.ReleaseProgram(m_program);
        m_program = 0;
    }

    if (m_commandQueue != 0) {
        m_cl.ReleaseCommandQueue(m_commandQueue);
        m_commandQueue = 0;
    }

    if (m_context != 0) {
        m_cl.ReleaseContext(m_context);
        m_context = 0;
    }

    m_cl.Close();
}

int GpuX::GetClError()
{
    return m_err;
}

int GpuX::Finish()
{
    m_err = m_cl.Finish(m_commandQueue);
    return m_err == CL_SUCCESS ? GpuX::ErrorCode::SUCCESS : GpuX::ErrorCode::FINISH_FAILED;
}

GpuX::Buffer *GpuX::CreateBuffer(void *ptr, size_t size)
{
    GpuX::Buffer *buffer = new GpuX::Buffer(&m_cl, m_context, m_commandQueue);
    int ret = buffer->Init(ptr, size);
    if (ret != GpuX::ErrorCode::SUCCESS) {
        return nullptr;
    }
    return buffer;
}

void GpuX::DestroyBuffer(GpuX::Buffer *buffer)
{
    delete buffer;
}

GpuX::~GpuX()
{
    Destroy();
}